import numpy as np
import matplotlib.pyplot as plt

# 定义 N 和 n
N = 1024
n = np.arange(N)

# 输入信号
x = [int(np.round(np.sin(i) * 1024)) + int(np.round(np.cos(i) * 1024)) * 1j for i in n]

# 频域抽取的基2FFT
loop_num = int(np.log2(N))
data = np.zeros((loop_num + 1, N), dtype=np.complex128)
data[0] = x

for i in range(loop_num):
    k = i + 1
    for p in range(2 ** i):
        for j in range(N // (2 ** k)):
            data[i + 1][j + p * (N // (2 ** i))] = data[i, j + p * (N // (2 ** i))] + data[
                i, j + N // (2 ** k) + p * (N // (2 ** i))]
            data[i + 1][j + N // (2 ** k) + p * (N // (2 ** i))] = (data[i, j + p * (N // (2 ** i))] - data[
                i, j + N // (2 ** k) + p * (N // (2 ** i))]) * np.exp(-1j * 2 * np.pi * j / (2 ** i))

# 计算位反转索引
def bit_reverse(k, N):
    reversed_k = 0
    for _ in range(int(np.log2(N))):
        reversed_k = (reversed_k << 1) | (k & 1)
        k >>= 1
    return reversed_k

# 输出倒序
fft_out = np.zeros_like(data[0, :])
for k in range(N):
    fft_out[bit_reverse(k, N)] = data[loop_num, k]

# 使用 NumPy 的 FFT 进行比较
xf = np.fft.fft(x)

# 绘制结果
plt.figure(figsize=(12, 6))
plt.subplot(2, 1, 1)
plt.plot(abs(xf), label='NumPy FFT')
plt.title('Magnitude of FFT using NumPy')
plt.legend()

plt.subplot(2, 1, 2)
plt.plot(abs(fft_out - xf), label='Difference')
plt.title('Difference between Custom FFT and NumPy FFT')
plt.legend()

plt.tight_layout()
plt.show()